# -*- coding: utf-8 -*-
# @Time    : 2020/12/4 10:29
# @Author  : DuJiabao
# @File    : ImageProcess.py
# @Description : 用于简单的图片处理
from math import ceil
import numpy as np
import cv2


def spNoise(image: np.ndarray, prob=0.1) -> np.ndarray:
    """
    产生椒盐噪声
    :param image: 原始图片
    :param prob: 椒盐噪声概率
    :return: 加工之后的图片
    """
    output = np.zeros(image.shape, np.uint8)
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            rdn = np.random.random()
            # 小于prob的，产生噪声
            if rdn < prob:
                output[i][j] = np.random.randint(0, 255)
            # 否则保持原状
            else:
                output[i][j] = image[i][j]
    return output


def synthesis(bg: np.ndarray, item: np.ndarray, coordinate: [tuple, list]) -> tuple:
    """
    功能：合成图片，在bg上，贴上item。
    :param bg: 背景图片
    :param item: 被贴的物品，最好是已经被抠出来，背景值都为0的图片
    :param coordinate: 被贴的中心位置坐标（y, x）
    :return: 新的img, boundingbox（[ymin, ymax, xmin, xmax]）
    """
    ih, iw, _ = item.shape
    h, w, _ = bg.shape
    # 蒙版，白色底，为了防止贴图溢出，在bg外围加一圈
    mask = np.ones((2 * ih + h, 2 * iw + w, 3), np.uint8) * 255
    # 粘贴坐标因此发生了偏移
    new_coor = [coordinate[0] + ih, coordinate[1] + iw]

    # item的蒙版，有东西的区域为黑色，没东西的区域为白色
    mask_item = (item == 0) * 255
    # 将蒙版贴到新的mask上， 贴物品的区域为黑色，其余为白色
    mask[new_coor[0] - int(ih / 2):new_coor[0] + ceil(ih / 2),
    new_coor[1] - int(iw / 2):new_coor[1] + ceil(iw / 2)] = mask_item

    # 新的背景，与mask的shape相同
    new_bg = np.zeros_like(mask, dtype=np.uint8)
    new_bg[ih:-ih, iw:-iw] = bg
    new_bg = cv2.bitwise_and(mask, new_bg)  # 将需要放置东西的区域，全体归0
    # 粘贴上去
    new_bg[new_coor[0] - int(ih / 2):new_coor[0] + ceil(ih / 2),
    new_coor[1] - int(iw / 2):new_coor[1] + ceil(iw / 2)] = \
        cv2.bitwise_or(item, new_bg[new_coor[0] - int(ih / 2):new_coor[0] + ceil(ih / 2),
                             new_coor[1] - int(iw / 2):new_coor[1] + ceil(iw / 2)])

    bbox = [max(0, coordinate[0] - int(ih / 2)),
            min(h, coordinate[0] + ceil(ih / 2)),
            max(0, coordinate[1] - int(iw / 2)),
            min(w, coordinate[1] + ceil(iw / 2))]
    return new_bg[ih:-ih, iw:-iw], bbox


def baffleEllipsis(item: np.ndarray) -> np.ndarray:
    """
    在原来图片的基础上，加一个挡板，被挡区域为黑色。
    :param item: 原始图片
    :return: 处理之后的图片
    """
    ih, iw, _ = item.shape
    ellipsis = cv2.ellipse(np.zeros_like(item), (int(iw / 2), int(ih / 2)), (int(iw / 2), int(ih / 2)),
                           np.random.randint(0, 360), 0, 360, (255, 255, 255), -1)
    item = cv2.bitwise_and(ellipsis, item)
    return item


def picRotate(image: np.ndarray, angle, borderValue=None) -> np.ndarray:
    """
    图片旋转，顺时针为正，逆时针为负
    :param borderValue: 空白处填充的颜色，可选择输入三通道的颜色
    :param image: 需要旋转的图片
    :param angle: 旋转的角度
    :return: 旋转后的图片
    """
    # 中心旋转
    # grab the dimensions of the image and then determine the
    # center
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)

    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    if borderValue is None:
        borderValue = tuple(image[30, 30, :].tolist())
    img = cv2.warpAffine(image, M, (nW, nH), borderValue=borderValue)
    # perform the actual rotation and return the image
    return img


def pic_resize(img: np.ndarray, min: int, max: int, random=1) -> np.ndarray:
    """
    图片随机resize
    :param img: 原始图片
    :param min: shape下限
    :param max: shape上限
    :param random: 长宽差别的随机范围
    :return: 处理之后的图片
    """
    size = np.random.randint(min, max)
    y = size + np.random.randint(0, random)
    x = size + np.random.randint(0, random)
    img = cv2.resize(img, (y, x))
    return img


def addBlur(img: np.ndarray, level: int, mode=0) -> np.ndarray:
    """
    添加模糊
    :param img: 原始图片
    :param level: 模糊级别
    :param mode: 模式（0：中值模糊，1：高斯模糊）
    :return: 处理之后的图片
    """
    if mode == 0:
        # 中值模糊
        return cv2.blur(img, (level * 2 + 1, level * 2 + 1))
    else:
        # 高斯模糊
        return cv2.GaussianBlur(img, (level * 2 + 1, level * 2 + 1), 1)


def brightnessAdjust(img: np.ndarray, alpha: float, beta: float) -> np.ndarray:
    """
    随机调节图片的对比度和亮度
    :param img: 原始图片
    :param alpha: α范围
    :param beta: β范围
    :return: 处理之后的图片
    """
    img = img.astype(np.int16) * (1 + np.random.uniform(-alpha, alpha)) + np.random.randint(-beta, beta + 1)
    return np.clip(img, 0, 255).astype(np.uint8)


def cropPic(img, xmin, ymin, xmax, ymax):
    return img[ymin:ymax, xmin:xmax]


if __name__ == '__main__':
    pass
